from prophet import Prophet
from common_import import *


def preprocess_data(data):
    # 将为0的异常值设置为None
    data["quantity"][data["quantity"] == 0] = None
    data["sales_price"][data["sales_price"] == 0] = None
    data["cost_price"][data["cost_price"] == 0] = None
    return data


def prophet_forecast(data, column_name):
    # 创建Prophet模型

    # 准备数据
    df = pd.DataFrame({"ds": data["date"], "y": data[column_name]})
    model = Prophet(
        yearly_seasonality=True,
        weekly_seasonality=True,
        daily_seasonality=False,  # 默认日规律不启用
    )

    # 训练模型
    model.fit(df)

    # 创建未来7天的时间数据
    future = model.make_future_dataframe(periods=7)

    # 进行预测
    forecast = model.predict(future)

    # 只保留预测的部分，并返回结果
    return forecast[["ds", "yhat"]].tail(7)


def forecast_future_sales(data):
    # 预处理数据
    data = preprocess_data(data)

    # 预测未来7天的销售量
    quantity_forecast = prophet_forecast(data, "quantity")

    # 预测未来7天的销售价格
    sales_price_forecast = prophet_forecast(data, "sales_price")

    # 预测未来7天的成本价格
    cost_price_forecast = prophet_forecast(data, "cost_price")

    # 将结果合并到一个结构化数组中
    result = np.array(
        list(
            zip(
                quantity_forecast["ds"],
                quantity_forecast["yhat"],
                sales_price_forecast["yhat"],
                cost_price_forecast["yhat"],
            )
        ),
        dtype=[
            ("date", "datetime64[D]"),
            ("quantity", "f4"),
            ("sales_price", "f4"),
            ("cost_price", "f4"),
        ],
    )
    result["quantity"] = np.abs(result["quantity"])
    result["quantity"][result["quantity"] < 2.5] += 2.5
    return result


if __name__ == "__main__":
    focus_product = mapping.focus_products
    for i in focus_product:
        data = tool.get_np(f"problem3/{i}.csv")
        result = forecast_future_sales(data)
        tool.get_csv(result, f"product_forcast/{i}.csv")
